import logging
from typing import List, Dict, Any, Optional
from langchain.schema import Document
from .document_processor import DocumentProcessor
from .reranker import Reranker
from models.llm_models import RAGModel
from models.langgraph_workflows import RAGWorkflow
from config import Config
from config.thresholds import RAGThresholds

logger = logging.getLogger(__name__)

class RAGSystemWorkflow:
    """使用LangGraph工作流的RAG系统：整合文档处理和问答功能 - 保持原有接口"""
    
    def __init__(self, vector_db_path: str = None, use_reranker: bool = True):
        # 保持与原有RAGSystem相同的初始化逻辑
        self.document_processor = DocumentProcessor(vector_db_path)
        self.rag_model = RAGModel()
        self.use_reranker = use_reranker
        self.reranker = Reranker() if use_reranker else None
        
        # 初始化LangGraph工作流
        self.workflow = RAGWorkflow(self)
    
    def upload_document(self, file_path: str, collection_name: str = "default") -> Dict[str, Any]:
        """上传并处理文档 - 保持原有接口"""
        try:
            logger.info(f"开始上传文档: {file_path}")
            
            # 保持原有的文档处理逻辑
            vector_store = self.document_processor.process_document(file_path, collection_name)
            
            # 获取集合信息
            collection_info = self.document_processor.get_collection_info(collection_name)
            
            result = {
                "status": "success",
                "message": f"文档 {file_path} 上传成功",
                "collection_info": collection_info,
                "file_path": file_path
            }
            
            logger.info(f"文档上传完成: {file_path}")
            return result
            
        except Exception as e:
            logger.error(f"文档上传失败: {str(e)}")
            return {
                "status": "error",
                "message": f"文档上传失败: {str(e)}",
                "file_path": file_path
            }
    
    def ask_question(self, question: str, collection_name: str = "default", top_k: int = None, rerank_threshold: float = RAGThresholds.DEFAULT_RERANK_THRESHOLD) -> Dict[str, Any]:
        """基于RAG回答问题 - 使用LangGraph工作流但保持原有接口"""
        try:
            logger.info(f"收到问题: {question}")
            
            # 使用LangGraph工作流，但保持原有的调用方式
            result = self.workflow.ask_question(
                question=question,
                collection_name=collection_name,
                top_k=top_k,
                rerank_threshold=rerank_threshold
            )
            
            logger.info(f"工作流回答问题完成: {question}")
            
            # 确保返回结果包含所有必需字段
            if result.get("status") == "error":
                return {
                    "status": "error",
                    "question": question,
                    "answer": result.get("message", "处理失败"),
                    "sources": [],
                    "context_length": 0,
                    "documents_retrieved": 0,
                    "documents_used": 0
                }
            
            return result
            
        except Exception as e:
            logger.error(f"回答问题失败: {str(e)}")
            return {
                "status": "error",
                "question": question,
                "answer": f"回答问题失败: {str(e)}",
                "sources": [],
                "context_length": 0,
                "documents_retrieved": 0,
                "documents_used": 0
            }
    
    def get_collection_info(self, collection_name: str = "default") -> Dict[str, Any]:
        """获取集合信息 - 保持原有接口"""
        try:
            return self.document_processor.get_collection_info(collection_name)
        except Exception as e:
            logger.error(f"获取集合信息失败: {str(e)}")
            return {
                "status": "error",
                "message": f"获取集合信息失败: {str(e)}"
            }
    
    def list_collections(self) -> List[str]:
        """列出所有集合 - 保持原有接口"""
        try:
            # 保持原有逻辑
            return ["default"]
        except Exception as e:
            logger.error(f"列出集合失败: {str(e)}")
            return []
    
    def delete_collection(self, collection_name: str) -> Dict[str, Any]:
        """删除集合 - 保持原有接口"""
        try:
            logger.info(f"删除集合: {collection_name}")
            
            return {
                "status": "success",
                "message": f"集合 {collection_name} 删除成功"
            }
            
        except Exception as e:
            logger.error(f"删除集合失败: {str(e)}")
            return {
                "status": "error",
                "message": f"删除集合失败: {str(e)}"
            }
    
    # 添加工作流特有的方法，用于调试和监控
    def get_workflow_info(self) -> Dict[str, Any]:
        """获取工作流信息"""
        return {
            "workflow_type": "LangGraph",
            "use_reranker": self.use_reranker,
            "workflow_steps": [
                "retrieve_documents",
                "rerank_documents",
                "build_context",
                "generate_answer",
                "extract_sources",
                "create_result"
            ]
        }
    
    def run_workflow_with_tracing(self, question: str, collection_name: str = "default", top_k: int = None, rerank_threshold: float = None) -> Dict[str, Any]:
        """运行工作流并返回详细的执行信息"""
        try:
            result = self.workflow.run(question, collection_name, top_k, rerank_threshold)
            
            return {
                **result,
                "workflow_info": self.get_workflow_info(),
                "execution_trace": True
            }
            
        except Exception as e:
            logger.error(f"工作流执行失败: {str(e)}")
            return {
                "error": f"工作流执行失败: {str(e)}",
                "workflow_info": self.get_workflow_info()
            }
    
    # 保持原有的辅助方法，用于工作流内部调用
    def _build_context(self, documents: List[Document]) -> str:
        """构建上下文 - 保持原有方法"""
        context_parts = []
        
        for i, doc in enumerate(documents, 1):
            source = doc.metadata.get("source", "未知来源")
            content = doc.page_content.strip()
            
            context_parts.append(f"文档 {i} (来源: {source}):\n{content}\n")
        
        return "\n".join(context_parts)
    
    def _extract_sources(self, documents: List[Document]) -> List[Dict[str, str]]:
        """提取源文档信息 - 保持原有方法"""
        sources = []
        
        for doc in documents:
            source_info = {
                "source": doc.metadata.get("source", "未知来源"),
                "file_name": doc.metadata.get("file_name", "未知文件"),
                "file_type": doc.metadata.get("file_type", "未知类型"),
                "content_preview": doc.page_content[:200] + "..." if len(doc.page_content) > 200 else doc.page_content
            }
            sources.append(source_info)
        
        return sources 